package developer.zyk.ability.chat.ability.rag.vector;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.lang.TypeReference;
import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import developer.zyk.ability.knowledge.dao.entity.ChatKnowledgeRecordDO;
import developer.zyk.ability.knowledge.dao.entity.ChatKnowledgeRecordDetailDO;
import developer.zyk.ability.knowledge.dao.mappers.ChatKnowledgeRecordDetailMapper;
import developer.zyk.ability.knowledge.dao.mappers.ChatKnowledgeRecordMapper;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.ai.vectorstore.SimpleVectorStoreContent;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;

import java.time.LocalDateTime;
import java.util.*;
import java.util.function.Function;
import java.util.stream.Collectors;

/**
 * @author zhangyongkang
 * @time 2025/5/13 17:22
 * @description 多租户版本的知识库
 */
@Slf4j
public class MysqlVectorStore extends SimpleVectorStore implements UserStoreVector {
    private final ChatKnowledgeRecordDetailMapper chatKnowledgeRecordDetailMapper;
    private final ChatKnowledgeRecordMapper chatKnowledgeRecordMapper;

    public MysqlVectorStore(SimpleVectorStoreBuilder builder,
                            ChatKnowledgeRecordDetailMapper detailMapper,
                            ChatKnowledgeRecordMapper chatKnowledgeRecordMapper
    ) {
        super(builder);
        this.chatKnowledgeRecordDetailMapper = detailMapper;
        this.chatKnowledgeRecordMapper = chatKnowledgeRecordMapper;
    }

    @Override
    public void doAdd(List<Document> documents) {
        super.doAdd(documents);
        List<ChatKnowledgeRecordDetailDO> recordDetails = documents.stream()
                .map(Document::getId)
                .map(a -> {
                    SimpleVectorStoreContent content = store.get(a);
                    ChatKnowledgeRecordDetailDO ChatKnowledgeRecordDetailDO = new ChatKnowledgeRecordDetailDO();
                    ChatKnowledgeRecordDetailDO.setRecordId((Long) content.getMetadata().get("recordId"));
                    ChatKnowledgeRecordDetailDO.setDocumentOutId(a);
                    ChatKnowledgeRecordDetailDO.setEmbeddingText(JSONUtil.toJsonStr(content.getEmbedding()));
                    ChatKnowledgeRecordDetailDO.setMetaData(JSONUtil.toJsonStr(content.getMetadata()));
                    ChatKnowledgeRecordDetailDO.setOriginText(content.getText());
                    ChatKnowledgeRecordDetailDO.setCreateTime(LocalDateTime.now());
                    return ChatKnowledgeRecordDetailDO;
                })
                .toList();
        chatKnowledgeRecordDetailMapper.insert(recordDetails);
    }

    @Override
    public VectorStoreObservationContext.Builder createObservationContextBuilder(String operationName) {
        return super.createObservationContextBuilder(operationName);
    }

    @Override
    public List<Document> doSimilaritySearch(SearchRequest request) {
        List<Document> documents = super.doSimilaritySearch(request);
        return documents;
    }

    @Override
    public void doDelete(List<String> idList) {
        super.doDelete(idList);
//        Optional.of(idList).orElse(new ArrayList<>())
//                .stream()
//                .filter(Objects::nonNull)
//                .forEach(a -> {
//                    Optional.ofNullable(store.get(a))
//                            .map(SimpleVectorStoreContent::getMetadata)
//                            .map(b -> b.get("recordId"))
//                            .map(b -> (Long) b)
//                            .ifPresent(recordId -> {
//                                chatUnstructuredRecordDetailMapper.delete(
//                                        new QueryWrapper<ChatKnowledgeRecordDetailDO>()
//                                                .lambda()
//                                                .eq(ChatKnowledgeRecordDetailDO::getRecordId, recordId)
//                                                .eq(ChatKnowledgeRecordDetailDO::getDocumentOutId, a)
//                                );
//                            });
//
//                });
    }


    @Override
    public void load(List<Long> recordIds) {
        List<ChatKnowledgeRecordDetailDO> details = chatKnowledgeRecordDetailMapper.selectList(
                new QueryWrapper<ChatKnowledgeRecordDetailDO>()
                        .lambda()
                        .in(ChatKnowledgeRecordDetailDO::getRecordId, recordIds)

        );
        doLoad(details);
    }

    @Override
    public void loadAllOnStart() {
        List<ChatKnowledgeRecordDetailDO> details = chatKnowledgeRecordDetailMapper.selectList(
                new QueryWrapper<ChatKnowledgeRecordDetailDO>()
                        .lambda()
        );
        doLoad(details);
    }

    @Override
    public void realDelete(List<Long> recordIds) {
        if (CollUtil.isEmpty(recordIds)) {
            return;
        }

        List<ChatKnowledgeRecordDetailDO> details = chatKnowledgeRecordDetailMapper.selectList(
                new QueryWrapper<ChatKnowledgeRecordDetailDO>()
                        .lambda()
                        .in(ChatKnowledgeRecordDetailDO::getRecordId, recordIds)
        );
        List<String> documentIds = details.stream()
                .filter(Objects::nonNull)
                .map(ChatKnowledgeRecordDetailDO::getDocumentOutId)
                .filter(Objects::nonNull)
                .toList();

        this.doDelete(documentIds);

        chatKnowledgeRecordDetailMapper.delete(
                new QueryWrapper<ChatKnowledgeRecordDetailDO>()
                        .lambda()
                        .in(ChatKnowledgeRecordDetailDO::getRecordId, recordIds)
        );
    }

    private void doLoad(
            List<ChatKnowledgeRecordDetailDO> recordDetails
    ) {
        //加载数据库的内容至内存store
        List<Long> recordIds = recordDetails.stream()
                .filter(Objects::nonNull)
                .map(ChatKnowledgeRecordDetailDO::getRecordId)
                .toList();
        if (CollUtil.isEmpty(recordIds)) {
            return;
        }

        Map<Long, ChatKnowledgeRecordDO> recordMap = chatKnowledgeRecordMapper.selectList(
                        new QueryWrapper<ChatKnowledgeRecordDO>()
                                .lambda()
                                .in(ChatKnowledgeRecordDO::getId, recordIds)
                )
                .stream()
                .collect(Collectors.toMap(ChatKnowledgeRecordDO::getId, Function.identity(), (a, b) -> b));

        recordDetails.stream()
                .filter(Objects::nonNull)
                .filter(a -> {
                    //没启用的就不加载了
                    return Optional.ofNullable(recordMap.get(a.getRecordId()))
                            .map(ChatKnowledgeRecordDO::getEnable)
                            .map(en -> Objects.equals(en, 1))
                            .orElse(false);
                })
                .forEach(a -> {
                    String documentOutId = a.getDocumentOutId();
                    double[] array = JSONUtil.toList(a.getEmbeddingText(), Float.class)
                            .stream()
                            .mapToDouble(b -> b)
                            .toArray();
                    float[] embedding = new float[array.length];
                    for (int i = 0; i < array.length; i++) {
                        embedding[i] = (float) array[i];
                    }
                    Map<String, Object> metadata =
                            Optional.ofNullable(a.getMetaData())
                                    .map(metaJson -> {
                                        return JSONUtil.toBean(metaJson,
                                                new TypeReference<Map<String, Object>>() {
                                                }, true
                                        );
                                    })
                                    .orElse(new HashMap<>());

                    store.put(documentOutId,
                            new SimpleVectorStoreContent(
                                    a.getOriginText(),
                                    metadata,
                                    embedding
                            ));

                });

    }
}
